import numpy as np
import matplotlib.pyplot as plt
import os
import matplotlib as mpl

def ensure_directory_exists(directory):
    """确保目录存在，如果不存在则创建"""
    if not os.path.exists(directory):
        os.makedirs(directory)
        print(f"创建目录: {directory}")

def set_plot_style():
    """设置绘图样式"""
    # 设置全局字体为黑体（如果支持）
    try:
        plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
        plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号
    except:
        plt.rcParams['font.sans-serif'] = ['DejaVu Sans']  # 备选字体
    
    # 设置全局样式
    plt.style.use('ggplot')
    mpl.rcParams['figure.figsize'] = (10, 8)
    mpl.rcParams['axes.grid'] = True
    mpl.rcParams['grid.alpha'] = 0.3
    mpl.rcParams['lines.linewidth'] = 2
    mpl.rcParams['axes.titlesize'] = 14
    mpl.rcParams['axes.labelsize'] = 12
    mpl.rcParams['xtick.labelsize'] = 10
    mpl.rcParams['ytick.labelsize'] = 10
    mpl.rcParams['legend.fontsize'] = 12
    mpl.rcParams['figure.titlesize'] = 16
    
    print("绘图样式已设置")

def calculate_errors(reference_trajectory, states_history):
    """
    计算轨迹跟踪误差
    reference_trajectory: 参考轨迹
    states_history: 状态历史
    返回包含各种误差指标的字典
    """
    # 确保长度匹配
    min_length = min(len(reference_trajectory), len(states_history))
    reference_trajectory = reference_trajectory[:min_length]
    states_history = states_history[:min_length]
    
    # 计算位置误差
    position_errors = reference_trajectory[:, :2] - states_history[:, :2]
    distance_errors = np.sqrt(np.sum(position_errors**2, axis=1))
    
    # 计算速度误差
    velocity_errors = reference_trajectory[:, 2:] - states_history[:, 2:]
    speed_errors = np.sqrt(np.sum(velocity_errors**2, axis=1))
    
    # 计算均方根误差
    rmse = np.sqrt(np.mean(distance_errors**2))
    vel_rmse = np.sqrt(np.mean(speed_errors**2))
    
    # 计算最大误差
    max_error = np.max(distance_errors)
    max_vel_error = np.max(speed_errors)
    
    # 计算平均误差
    mean_error = np.mean(distance_errors)
    mean_vel_error = np.mean(speed_errors)
    
    return {
        'position_errors': position_errors,
        'distance_errors': distance_errors,
        'velocity_errors': velocity_errors,
        'speed_errors': speed_errors,
        'rmse': rmse,
        'vel_rmse': vel_rmse,
        'max_error': max_error,
        'max_vel_error': max_vel_error,
        'mean_error': mean_error,
        'mean_vel_error': mean_vel_error
    }

def save_simulation_data(results, save_path):
    """
    保存仿真数据
    results: 仿真结果
    save_path: 保存路径
    """
    np.savez(
        save_path,
        time_points=results['time_points'],
        states_history=results['states_history'],
        reference_trajectory=results['reference_trajectory'],
        controls_history=results['controls_history'] if 'controls_history' in results else None
    )
    print(f"仿真数据已保存到: {save_path}")

def load_simulation_data(load_path):
    """
    加载仿真数据
    load_path: 加载路径
    返回仿真结果字典
    """
    data = np.load(load_path)
    results = {
        'time_points': data['time_points'],
        'states_history': data['states_history'],
        'reference_trajectory': data['reference_trajectory'],
    }
    
    # 如果有控制输入历史，也加载
    if 'controls_history' in data:
        results['controls_history'] = data['controls_history']
        
    print(f"仿真数据已从 {load_path} 加载")
    return results

def print_simulation_summary(errors):
    """
    打印仿真摘要
    errors: 误差字典
    """
    print("\n" + "=" * 30)
    print("仿真结果摘要")
    print("=" * 30)
    print(f"位置RMSE: {errors['rmse']:.4f}")
    print(f"速度RMSE: {errors['vel_rmse']:.4f}")
    print(f"最大位置误差: {errors['max_error']:.4f}")
    print(f"最大速度误差: {errors['max_vel_error']:.4f}")
    print(f"平均位置误差: {errors['mean_error']:.4f}")
    print(f"平均速度误差: {errors['mean_vel_error']:.4f}")
    print("=" * 30)

def compute_control_effort(control_inputs):
    """
    计算控制代价
    control_inputs: 控制输入历史
    """
    # 计算每个时间步的控制能量 (u^T * u)
    control_energy = np.sum(control_inputs**2, axis=1)
    
    # 总控制代价
    total_control_effort = np.sum(control_energy)
    
    # 平均控制代价
    avg_control_effort = np.mean(control_energy)
    
    return {
        'control_energy': control_energy,
        'total_control_effort': total_control_effort,
        'avg_control_effort': avg_control_effort
    }

def compare_trajectories(reference_trajectory, actual_trajectory, time_points, save_path=None):
    """
    比较参考轨迹和实际轨迹
    reference_trajectory: 参考轨迹
    actual_trajectory: 实际轨迹
    time_points: 时间点
    save_path: 保存路径
    """
    # 确保长度一致
    min_len = min(len(reference_trajectory), len(actual_trajectory))
    reference_trajectory = reference_trajectory[:min_len]
    actual_trajectory = actual_trajectory[:min_len]
    time_points = time_points[:min_len]
    
    # 计算位置误差
    position_error = reference_trajectory[:, :2] - actual_trajectory[:, :2]
    error_distance = np.sqrt(np.sum(position_error**2, axis=1))
    
    # 创建图形
    fig, axes = plt.subplots(2, 1, figsize=(10, 12))
    
    # 绘制XY轨迹
    axes[0].plot(reference_trajectory[:, 0], reference_trajectory[:, 1], 'b-', label='参考轨迹')
    axes[0].plot(actual_trajectory[:, 0], actual_trajectory[:, 1], 'r--', label='实际轨迹')
    axes[0].set_xlabel('X位置')
    axes[0].set_ylabel('Y位置')
    axes[0].set_title('轨迹比较')
    axes[0].grid(True)
    axes[0].legend()
    axes[0].axis('equal')
    
    # 绘制跟踪误差
    axes[1].plot(time_points, error_distance, 'k-')
    axes[1].set_xlabel('时间(s)')
    axes[1].set_ylabel('跟踪误差')
    axes[1].set_title('跟踪误差随时间变化')
    axes[1].grid(True)
    
    # 调整布局
    plt.tight_layout()
    
    # 保存图像
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
    
    plt.show() 